#!/usr/bin/env python3
# -*- coding:utf-8 -*-
# @Author: renjin@bit.edu.cn

import math
import cv2
import base64
import copy
import math
import argparse
import os
import numpy as np
from datetime import datetime
from spirems import SMSBagPlayer
from spirems.image_io.visual_helper import draw_charts_v3, load_a2rl_logo, track_boundary_parse, draw_track_map_v2
from spirems import Publisher, Subscriber, def_msg, cvimg2sms


def smsbag_racecar_replay(smsbag_path, output_video_dir=''):
    url_img = "/sensor/camera/camera_fl/compressed"
    url_ego_loc = "/flyeagle/a2rl/observer/ego_loc"
    url_ego_state = "/flyeagle/a2rl/observer/ego_state"
    url_ctl_status = "/flyeagle/a2rl/controller/status"
    url_ctl_debug = "/flyeagle/a2rl/controller/debug"

    player = SMSBagPlayer(smsbag_path)
    t0 = 0
    nc = 0
    visual_msg = def_msg('std_msgs::Null')
    visual_msg['timestamp'] = 0.0
    visual_msg['ego_position'] = [0, 0, 0]
    visual_msg['ego_orientation_ypr'] = [0, 0, 0]
    visual_msg['ego_velocity'] = [0, 0, 0]
    visual_msg['ego_acceleration'] = [0, 0, 0]
    visual_msg['ice_actual_gear'] = 1
    visual_msg['ice_actual_throttle'] = 0.0
    visual_msg['ice_engine_speed_rpm'] = 0.0
    visual_msg['ice_water_temp_deg_c'] = 0.0
    visual_msg['ice_oil_temp_deg_c'] = 0.0
    visual_msg['brake_f'] = 0.0
    visual_msg['brake_r'] = 0.0
    visual_msg['slip_f'] = 0.0
    visual_msg['slip_r'] = 0.0
    visual_msg['steering'] = 0.0
    visual_msg['tyre_temp_fl'] = [0, 0, 0]
    visual_msg['tyre_temp_fr'] = [0, 0, 0]
    visual_msg['tyre_temp_rl'] = [0, 0, 0]
    visual_msg['tyre_temp_rr'] = [0, 0, 0]
    visual_msg['brake_disk_temp_fl'] = 0.0
    visual_msg['brake_disk_temp_fr'] = 0.0
    visual_msg['brake_disk_temp_rl'] = 0.0
    visual_msg['brake_disk_temp_rr'] = 0.0
    visual_msg['error_lateral'] = 0.0
    visual_msg['error_yaw'] = 0.0
    visual_msg['error_speed'] = 0.0
    visual_msg['lap_count'] = 1

    if len(output_video_dir) > 0:
        fourcc = cv2.VideoWriter_fourcc(*'mp4v')  # MP4格式编码器
        output_video = cv2.VideoWriter(
            os.path.join(output_video_dir, datetime.now().strftime("smsreplay_%Y-%m-%d_%H-%M-%S.mp4")),  # 输出文件名
            fourcc,
            60,
            (1280, 720)  # 必须与帧尺寸一致
        )

    while 1:
        msg = player.next()
        if msg is None:
            break
        url = msg['msg']['url']
        tim = msg['msg']['timerec']
        if url == url_img:
            img_raw = msg['raw']
            if t0 == 0:
                t0 = tim
        elif url == url_ego_state:
            visual_msg['ego_acceleration'] = [
                msg['msg']['acceleration']['x'], msg['msg']['acceleration']['y'], msg['msg']['acceleration']['z']
            ]
            visual_msg['ego_velocity'] = [
                msg['msg']['velocity']['x'], msg['msg']['velocity']['y'], msg['msg']['velocity']['z']
            ]
        elif url == url_ego_loc:
            visual_msg['ego_position'] = [
                msg['msg']['position']['x'], msg['msg']['position']['y'], msg['msg']['position']['z']
            ]
        elif url == url_ctl_status:
            visual_msg['ice_actual_gear'] = msg['msg']['gear']
            visual_msg['ice_actual_throttle'] = msg['msg']['throttle']
            visual_msg['steering'] = msg['msg']['steering']
            visual_msg['brake_f'] = msg['msg']['front_brake']
            visual_msg['brake_r'] = msg['msg']['rear_brake']
            visual_msg['slip_f'] = msg['msg']['slip_f']
            visual_msg['slip_r'] = msg['msg']['slip_r']
        elif url == url_ctl_debug:
            visual_msg['error_lateral'] = msg['msg']['lateral_error']
            visual_msg['error_yaw'] = msg['msg']['yaw_error']
            visual_msg['error_speed'] = msg['msg']['speed_error']

        if t0 > 0:
            n_img = int((tim - t0) * 1000 / 16.666667)
            if n_img != nc:
                nc = n_img
                visual_msg['timestamp'] = tim
                img_show = draw_charts_v3(img_raw, visual_msg)

                if len(output_video_dir) > 0 and img_show.shape[0] == 720 and img_show.shape[1] == 1280:
                    output_video.write(img_show)

                cv2.imshow("img_raw", img_show)
                cv2.waitKey(5)

    if len(output_video_dir) > 0:
        output_video.release()

def main():
    parser = argparse.ArgumentParser(description="ROS2Bag A2RL RaceCar 回放程序")
    parser.add_argument(
        '-s', '--sms-bag',
        type=str,
        default='/home/amov/Pictures/smsbag_2025-10-18_20-01-42',
        help='输入的smsbag路径'
    )
    parser.add_argument(
        '-o', '--output-dir',
        type=str,
        default='',
        help='输出视频文件夹'
    )
    args = parser.parse_args()

    print("Input SMSBag: {}".format(args.sms_bag))
    if not os.path.isdir(args.sms_bag):
        print("SMSBag不存在")
        exit(1)

    smsbag_path = args.sms_bag
    smsbag_racecar_replay(smsbag_path, args.output_dir)


if __name__ == "__main__":
    main()
